package com.ht.module.ai.config;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ht.constant.AiChatMessageRoleEnum;
import com.ht.module.ai.entity.AiChatMessage;
import com.ht.module.ai.entity.AiChatSession;
import com.ht.module.ai.mapper.AiChatMessageMapper;
import com.ht.module.ai.mapper.AiChatSessionMapper;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.data.message.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.Builder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.Resource;
import java.time.LocalDateTime;
import java.util.*;

/**
 * @author hejialun
 * @ClassName PersistentChatMemoryStore
 * @description: TODO（ai消息记忆）
 * @Date: 2025/6/9 15:25
 */
@Slf4j
@Configuration
public class PersistentChatMemoryStore implements ChatMemoryStore {
    @Resource
    private AiChatMessageMapper messageMapper;


    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String sessionId = (String) memoryId;
        List<ChatMessage> messages = new ArrayList<>();
        if (StrUtil.isEmpty(sessionId)) {
            return messages;
        }

        List<AiChatMessage> dbMessages = messageMapper.selectList(
                Wrappers.<AiChatMessage>lambdaQuery()
                        .eq(AiChatMessage::getSessionId, sessionId)
                        .orderByAsc(AiChatMessage::getMessageTime)
        );

        // 转换为 LangChain4j 的 ChatMessage
        for (AiChatMessage msg : dbMessages) {
            if (StrUtil.isEmpty(msg.getContent())) {
                msg.setContent("");
            }
            switch (msg.getRole()) {
                case "USER":
                    messages.add(new UserMessage(msg.getContent()));
                    break;
                case "AI":
                    messages.add(new AiMessage(msg.getContent()));
                    break;
                case "AI_CONFIRM_TOOL":
                    List<Map<String, String>> list = new ArrayList<>();
                    try {
                        list = new ObjectMapper().readValue(msg.getContent(), new TypeReference<>() {});
                    } catch (JsonProcessingException e) {
                       log.error("数据转换异常！");
                    }
                    List<ToolExecutionRequest> toolExecutionRequests=new ArrayList<>();
                    for (Map<String, String> map : list) {
                        ToolExecutionRequest request = ToolExecutionRequest.builder()
                                .id(map.get("id"))
                                .name(map.get("name"))
                                .arguments(map.get("arguments"))
                                .build();
                        toolExecutionRequests.add(request);
                    }
                    AiMessage aiMessage = new AiMessage(toolExecutionRequests);
                    messages.add(aiMessage);
                    break;
                case "TOOL_RESULT":
                    Map<String, String> map = new HashMap<>();
                    try {
                        map = new ObjectMapper().readValue(msg.getContent(), new TypeReference<Map<String, String>>() {});
                    } catch (JsonProcessingException e) {
                        log.error("数据转换异常！");
                    }
                    messages.add(new ToolExecutionResultMessage(
                            map.get("id"),
                            map.get("toolName"),
                            map.get("text")
                    ));
                    break;
                case "SYSTEM":
                    messages.add(new SystemMessage(msg.getContent()));
                    break;
                default:
                    // 忽略其他角色
                    break;
            }
        }
        return messages;
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String sessionId = (String) memoryId;
        if (StrUtil.isEmpty(sessionId) || CollUtil.isEmpty(messages)) {
            return;
        }
        ObjectMapper objectMapper = new ObjectMapper();
        // 只保存最新的消息，避免重复
        ChatMessage lastMessage = messages.get(messages.size() - 1);

        AiChatMessage aiMsg = new AiChatMessage();
        aiMsg.setSessionId(sessionId);
        aiMsg.setMessageTime(LocalDateTime.now());

        //只更新用户的消息
        if (lastMessage instanceof UserMessage msg) {
            //用户消息
            aiMsg.setRole(AiChatMessageRoleEnum.USER.getCode());
            aiMsg.setContent(msg.singleText());
        }
        if (lastMessage instanceof UserMessage msg) {
            //用户消息
            aiMsg.setRole(AiChatMessageRoleEnum.USER.getCode());
            aiMsg.setContent(msg.singleText());
        } else if (lastMessage instanceof AiMessage msg) {
            //ai消息
            List<ToolExecutionRequest> toolExecutionRequests = msg.toolExecutionRequests();

            if (CollUtil.isEmpty(toolExecutionRequests)) {
                aiMsg.setRole(AiChatMessageRoleEnum.AI.getCode());
                //ai默认消息
                aiMsg.setContent(msg.text());
            } else {
                //ai工具确认消息
                aiMsg.setRole(AiChatMessageRoleEnum.AI_CONFIRM_TOOL.getCode());
                List<Map<String, Object>> serializedList = new ArrayList<>();
                for (ToolExecutionRequest req : toolExecutionRequests) {
                    Map<String, Object> map = new HashMap<>();
                    map.put("id", req.id());
                    map.put("name", req.name());
                    map.put("arguments", req.arguments()); // arguments 是 StructuredToolInput
                    serializedList.add(map);
                }

                String json = null;
                try {
                    json = objectMapper.writeValueAsString(serializedList);
                } catch (JsonProcessingException e) {
                    log.error("工具类返回处理异常：" + e.getMessage());
                }
                aiMsg.setContent(json);
            }
        } else if (lastMessage instanceof ToolExecutionResultMessage msg) {
            //工具消息
            aiMsg.setRole(AiChatMessageRoleEnum.TOOL_RESULT.getCode());
            Map<String, Object> map = new HashMap<>();
            map.put("id", msg.id());
            map.put("text", msg.text());
            map.put("toolName", msg.toolName());
            String json = null;
            try {
                json = objectMapper.writeValueAsString(map);
            } catch (JsonProcessingException e) {
                log.error("工具类返回处理异常：" + e.getMessage());
            }
            aiMsg.setContent(json);
        }
        if (StrUtil.isNotEmpty(aiMsg.getContent())) {
            messageMapper.insert(aiMsg);
        } else if (lastMessage instanceof SystemMessage msg) {
            //系统消息
            aiMsg.setRole(AiChatMessageRoleEnum.SYSTEM.getCode());
            aiMsg.setContent(msg.text());
        }


    }

    @Override
    public void deleteMessages(Object memoryId) {
        String sessionId = (String) memoryId;
        if (StrUtil.isNotEmpty(sessionId)) {
            messageMapper.delete(
                    Wrappers.<AiChatMessage>lambdaQuery().eq(AiChatMessage::getSessionId, sessionId)
            );
        }
    }
}
